# mcp_client.py
from fastmcp import Client
import asyncio
import logging
import mcp

# 配置日志记录
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("calculator_mcp_client")

class AIClient:
    def __init__(self, host="localhost", port=25565):
        self.host = host
        self.port = port
        self.client = None
        self.connected = False

    async def connect(self):
        try:
            # 构建服务器连接URL
            transport_url = f"http://{self.host}:{self.port}/mcp"
            self.client = Client(transport=transport_url)
            
            print(f"正在连接到计算服务器 {self.host}:{self.port}...")
            
            # 连接服务器（FastMCP客户端会自动连接）
            # 不需要手动调用_connect()，而是通过上下文管理器或直接使用client实例
            self.connected = True
            print("已连接到计算服务器！")
            return True
            
        except Exception as e:
            logger.error(f"连接服务器失败: {e}")
            print(f"连接错误: {e}")
            self.connected = False
            return False
    
    async def call_calculator(self, tool_name: str, *args):
        """调用服务器端的工具函数，使用fastmcp的API"""
        if not self.client or not self.connected:
            raise RuntimeError("客户端未连接或已停止")
            
        # 构建工具调用参数，与server.py中的工具参数匹配
        if len(args) >= 2:
            arguments = {"a": args[0], "b": args[1]}
        else:
            raise ValueError("需要2个参数")
            
        logger.info(f"调用工具: {tool_name}，参数: {arguments}")
        
        try:
            # 使用async with确保调用过程中的连接稳定性
            async with self.client:
                # 使用fastmcp的call_tool方法调用服务端工具
                result = await self.client.call_tool(
                    tool_name,
                    arguments=arguments
                )
                return result
        except Exception as e:
            logger.error(f"调用工具失败: {e}")
            raise

    async def cleanup(self):
        """清理资源"""
        self.connected = False
        
    async def disconnect(self):
        """断开连接"""
        self.connected = False
        # FastMCP的Client会自动管理连接的关闭
        print("已断开与服务器的连接")
        await self.client._disconnect()
    async def get_tools(self) -> list[mcp.types.Tool]:
        """获取服务器端注册的所有工具"""
        if not self.client or not self.connected:
            raise RuntimeError("客户端未连接或已停止")
            
        try:
            async with self.client:
                tools = await self.client.list_tools()
                return tools
        except Exception as e:
            logger.error(f"获取工具列表失败: {e}")
            raise